import streamlit as st
import math


def calculate_memory(
        non_moe_layers: int,
        moe_layers: int,
        hidden_size: int,
        compress_dim: int,
        activated_experts: int,
        params: float,
        seq_len: int,
        batch_size: int,
        concurrency: int,
        dtype_bytes: int = 2
) -> dict:
    """DeepSeek MoE显存计算核心逻辑"""
    # 模型权重显存
    weight_mem_gb = params * 1e9 * dtype_bytes / (1024 ** 3)

    # KV缓存计算
    kv_non_moe = 2 * batch_size * seq_len * hidden_size * dtype_bytes * non_moe_layers / 1e9
    kv_moe = 2 * batch_size * seq_len * activated_experts * compress_dim * dtype_bytes * moe_layers / 1e9

    # 激活值显存
    activation_non_moe = batch_size * seq_len * hidden_size * dtype_bytes * non_moe_layers / 1e9
    activation_moe = batch_size * seq_len * activated_experts * compress_dim * dtype_bytes * moe_layers / 1e9

    # 单请求总临时显存
    per_request_mem = kv_non_moe + kv_moe + activation_non_moe + activation_moe

    # 并发总显存
    total_mem = weight_mem_gb + per_request_mem * concurrency

    return {
        "weight_mem": weight_mem_gb,
        "kv_non_moe": kv_non_moe,
        "kv_moe": kv_moe,
        "activation_non_moe": activation_non_moe,
        "activation_moe": activation_moe,
        "total_mem": total_mem
    }


# Streamlit可视化界面
st.title("DeepSeek MoE显存计算器")
st.caption("根据牛山AI公园评估表设计 v1.1")

with st.sidebar:
    st.header("模型参数配置")
    params = st.number_input("模型参数量（B）", value=37.0)
    non_moe_layers = st.number_input("非MoE层数", value=3)
    moe_layers = st.number_input("MoE层数", value=32)
    hidden_size = st.number_input("隐藏层维度", value=7168)
    compress_dim = st.number_input("压缩维度", value=512)
    activated_experts = st.number_input("激活专家数", value=8)

col1, col2 = st.columns(2)
with col1:
    st.header("推理参数")
    seq_len = st.slider("序列长度", 256, 4096, 2000)
    batch_size = st.selectbox("批处理大小", [1, 2, 4, 8, 16], index=0)
    concurrency = st.slider("并发用户数", 1, 100, 20)

with col2:
    st.header("精度设置")
    dtype = st.selectbox("数据类型", ["FP16", "FP32", "BF16"], index=0)
    dtype_bytes = 2 if dtype != "FP32" else 4

# 执行计算
result = calculate_memory(
    non_moe_layers=non_moe_layers,
    moe_layers=moe_layers,
    hidden_size=hidden_size,
    compress_dim=compress_dim,
    activated_experts=activated_experts,
    params=params,
    seq_len=seq_len,
    batch_size=batch_size,
    concurrency=concurrency,
    dtype_bytes=dtype_bytes
)

# 可视化结果
st.divider()
st.subheader("显存分析报告")

metric_cols = st.columns(3)
with metric_cols[0]:
    st.metric("模型权重显存", f"{result['weight_mem']:.2f} GB")
with metric_cols[1]:
    st.metric("单请求动态显存",
              f"{(result['kv_non_moe'] + result['kv_moe'] + result['activation_non_moe'] + result['activation_moe']):.2f} GB")
with metric_cols[2]:
    st.metric("总显存需求", f"{result['total_mem']:.2f} GB")

# 详细分析
with st.expander("详细组件分析"):
    components = [
        ("KV缓存-非MoE层", result['kv_non_moe']),
        ("KV缓存-MoE层", result['kv_moe']),
        ("激活值-非MoE层", result['activation_non_moe']),
        ("激活值-MoE层", result['activation_moe'])
    ]
    for label, value in components:
        st.progress(value / 2, f"{label}: {value:.2f} GB")

# 运行说明
st.divider()
st.write("""
**使用说明：**
1. 左侧栏配置模型结构参数
2. 右侧调整推理参数和精度
3. 结果自动更新，支持多设备对比
""")